/*!
  @file motion_solver_ransac.cpp
  @copyright 2013 Kubota Lab. All rights resereved.
*/

#include "mrover_vo/motion_solver_ransac.h"

#include <algorithm>
#include <numeric>
#include <limits>

#include <ros/ros.h>
#include <boost/random/mersenne_twister.hpp>
#include <boost/random/uniform_int.hpp>

namespace vo {

template <typename ValType>
void MotionSolverRANSAC<ValType>::initializeData() {
  input_pts .resize(3, matches_->size());
  target_pts.resize(3, matches_->size());
  for (size_t i = 0; i < matches_->size(); ++i) {
    cv::Point3f _pt_in = input_ ->sp->at(matches_->at(i).queryIdx);
    cv::Point3f _pt_tr = target_->sp->at(matches_->at(i).trainIdx);
    input_pts.col(i)  << _pt_in.x, _pt_in.y, _pt_in.z;
    target_pts.col(i) << _pt_tr.x, _pt_tr.y, _pt_tr.z;
  }

  method_->setInputFeature(input_);
  method_->setTargetFeature(target_);
  method_->setInputMatches(matches_);
}

template <typename ValType>
void MotionSolverRANSAC<ValType>::compute() {
  if (!initCompute()) {
    ROS_ERROR("error in two point solver");
    return;
  }

  ROS_INFO_STREAM("input " << input_->sp->size());
  ROS_INFO_STREAM("target " << target_->sp->size());
  ROS_INFO_STREAM("matches " << matches_->size());

  inliers_->clear();

  initializeData();
  int thresh_inlier = thresh_inlier_ratio_ * matches_->size();

  std::vector<int> num_inliers;
  Poses candidates;

  // start iteration
  for (int p = 0; p < num_itr_; ++p) {
    // compute poses
    PosesPtr hypotheses(new Poses);
    while (hypotheses->size() == 0)
      computePoseWithRandomSampling(hypotheses);

    // evaluate poses
    for (PosesItr h = hypotheses->begin(); h != hypotheses->end(); ++h) {
      std::vector<ValType> errors;
      computeReprojectionError(*h, errors);

      int n = 0;
      for (size_t i = 0; i < errors.size(); ++i) 
        n += (errors[i] < thresh_dist_)? 1: 0;
      if (n < thresh_inlier) continue;

      num_inliers.push_back(n);
      candidates.push_back(*h);

      //ROS_INFO_STREAM("Num of inliers(" << p << "): " << num_inliers.back());
    }
  }

  if (num_inliers.size() == 0) {
    ROS_ERROR("No poses found. Try to relax the threshold");
    return;
  }

  // choose the hypothesis with max inliers.
  int good_idx = std::distance(num_inliers.begin(), std::max_element(num_inliers.begin(), num_inliers.end()));
  ROS_INFO_STREAM(good_idx << " / " << num_inliers.size());
  convertEigen2CV<float>(candidates[good_idx].matrix(), pose_);

  std::vector<ValType> errors;
  computeReprojectionError(candidates[good_idx], errors);
  for (size_t i = 0; i < errors.size(); ++i) 
    if (errors[i] < thresh_dist_) inliers_->push_back(i);
  ROS_INFO_STREAM("Num of inliers (final): " << inliers_->size());
}

#if 0
template <typename ValType>
void MotionSolverRANSAC<ValType>::computePose2PT(const std::vector<int> &samples, Poses &iTt) {
  // Step 0. Convert the matrices
  MatrixXf X(3, num_sample);
  MatrixXf Y(3, num_sample);
  for (size_t i = 0; i < num_sample; ++i) {
    cv::Point3f _pt_in = input_sp_-> at(matches_->at(samples[i]).queryIdx);
    cv::Point3f _pt_tr = target_sp_->at(matches_->at(samples[i]).trainIdx);
    X.col(i) << _pt_tr.x, _pt_tr.y, _pt_tr.z;
    Y.col(i) << _pt_in.x, _pt_in.y, _pt_in.z;
  }

  // Step 1. Aligning axis
  Eigen::Vector3f z;  
  z << 0, 0, 1;

}
#endif

template <typename ValType>
void MotionSolverRANSAC<ValType>::computePoseWithRandomSampling(PosesPtr &poses) {
  std::vector<int> samples;
  generateRandomIndices(samples, matches_->size());
  method_->compute(samples, poses);
}

template <typename ValType>
ValType MotionSolverRANSAC<ValType>::computeScore(const Affine3T &iTt) {
  std::vector<ValType> errors;
  computeReprojectionError(iTt, errors);
  return std::accumulate(errors.begin(), errors.end(), 0);
}

template <typename ValType>
void MotionSolverRANSAC<ValType>::computeReprojectionError(const Affine3T &iTt, std::vector<ValType> &errors) {
  size_t num_points = matches_->size();

  const ValType POINT_GOES_BEHIND = std::numeric_limits<ValType>::max();

  MatrixXT transformed_pts = iTt.linear() * target_pts + iTt.translation().rowwise().replicate(num_points);
  MatrixXT projected_pts = K * transformed_pts;
  // @todo write this in better way?
  for (size_t i = 0; i < num_points; ++i) {
    projected_pts(0, i) /= projected_pts(2, i);
    projected_pts(1, i) /= projected_pts(2, i);
  }

  errors.resize(num_points);
  for (size_t i = 0; i < num_points; ++i) {
    if (transformed_pts(2, i) < 0) {
      errors[i] = POINT_GOES_BEHIND;
      continue;
    }

    ValType du = projected_pts(0, i) - input_->kp_vec[0]->at(matches_->at(i).queryIdx).pt.x;
    ValType dv = projected_pts(1, i) - input_->kp_vec[0]->at(matches_->at(i).queryIdx).pt.y;
    errors[i] = std::sqrt(du*du + dv*dv);
  }
}

template <typename ValType>
void MotionSolverRANSAC<ValType>::generateRandomIndices(std::vector<int> &indices, const size_t num_max) {
  size_t num_indices = method_->getRequiredSampleSize();
  indices.clear();

  static boost::mt19937 generator;
  boost::uniform_int<> dist(0, num_max - 1);

  while (indices.size() != num_indices) {
    int idx = dist(generator);
    int p = indices.size();
    std::vector<int>::iterator it = find(indices.begin(), indices.begin() + p, idx);
    if (it == indices.begin() + p && method_->checkGoodness(indices, idx)) {
      indices.push_back(idx);
    }
  }
}

}  // namespace vo
